"""
模型评估
"""
import numpy as np
import torch

from dataset import train_dataloader
from seq2seq.seq2seq_model import Seq2seq
import seq2seq.config as config
from torch.optim import Adam
import torch.nn.functional as f
from tqdm import tqdm

# 准备测试数据
test_data = [str(i) for i in np.random.randint(1, 1e8, size=[10])]
print(test_data)
test_data = sorted(test_data, key=lambda x: len(x), reverse=True)
test_data_length = torch.LongTensor([len(i) for i in test_data]).to(config.device)
input = torch.LongTensor(
    [config.num_sequence.transform(list(i), max_len=config.max_len, add_eos=True) for i in test_data]).to(config.device)

seq2seq = Seq2seq().to(config.device)
seq2seq.load_state_dict(torch.load(config.model_save_path))

indices = seq2seq.evaluation(input, test_data_length)
indices = np.array(indices).transpose()
# 反序列化观察结果
# print([config.num_sequence.inverse_transform(i) for i in indices])

result = []
for line in indices:
    print(line)
    temp_result = config.num_sequence.inverse_transform(line)
    cur_line = ""
    for word in temp_result:
        if word == config.num_sequence.EOS_TAG:
            break
        cur_line += word
    result.append(cur_line)
print(result)

target = [i + "0" for i in test_data]  # 我已经知道正确答案是最后加0
# 查看正确率
print(sum(i == j for i, j in zip(target, result)) / len(target))
